from __future__ import with_statement
import datetime
from Queue import Queue
import optparse
import os
import re
import socket
import sys
import threading
import time

from lib.classloader import loadpoc, autoLoadPocThread
import lib.classloader
from lib.data import banner
from lib.data import pocs
from lib.log import log, logToScreen



nmap_result_queue = Queue(100);
pro_semp = threading.Semaphore(5);
con_semp = threading.Semaphore(0);

logfile = None;
logfile_service = None;
targets = set();
targets_lock = threading.Lock();
scanned_targets = set()
scanned_targets_lock = threading.Lock();
round = 0;

class produceThread(threading.Thread):
    def __init__(self, prosemp, ip, consemp, callback):
        self.prosemp = prosemp;
        self.consemp = consemp
        self.ip = ip;
        self.callback = callback;
        super(produceThread, self).__init__();

    def run(self):
        global nmap_result_queue;
        global logfile_service;
        goods = {};
        goods["services"] = [];
        cmd = "nmap -sV --version-all -p1-65535 --open -n --host-timeout=3600s %s -T4"%(self.ip);
        logToScreen("info", cmd) ;
        result = os.popen(cmd).readlines();
        start = 0;
        for line in result:
            if line.find('Nmap scan') == 0:
                start = 1;
                goods["ip"] = self.ip;
                continue;
            if start == 1 and re.search('^[0-9]+/tcp', line):
                output=line.split();
                output[0] = output[0].replace("/tcp", "");
                goods["services"].append(
                    [output[0], output[1], output[2], ' '.join(output[3:])]);
        if len(goods["services"]) != 0 and len(goods["services"]) < 1000:#to prevent nmap generate 1-65535 is open
            nmap_result_queue.put(goods);
            self.callback();
            self.consemp.release();
            for item in goods["services"]:
                log(logfile_service, "\t".join([goods["ip"]]+ item));            
        self.prosemp.release();

class consumerThread(threading.Thread):
    def __init__(self, callback, semp):
        self.callback = callback;
        self.semp =semp;
        super(consumerThread, self).__init__();

    def run(self):
        global nmap_result_queue;
        self.semp.acquire();
        good = nmap_result_queue.get();

        if good and len(good.get("services")) != 0:
            for service in good["services"]:
                passvar = {};
                passvar["ip"] = good["ip"];
                passvar["port"] = service[0];
                passvar["service_type"] = service[2];
                passvar["service_version"] = service[3];
                for pockey in pocs.keys():
                    localpassvar = passvar.copy();
                    poc = pocs.get(pockey);
                    result = poc.execute(localpassvar);
                    result["poc"] = pockey;
                    self.callback(result);
        nmap_result_queue.task_done();

def reporter(result):
    global logfile;
    if result.get("success") and result["success"] == True:
        log(logfile, "\t".join([result["ip"], result["port"], result["service_version"], result["poc"], result["extra"] if result.get("extra") else ""]) + "\n")



def read_targets(target_file):
    tagetfile = open(target_file);
    targets = set();
    for line in tagetfile.readlines():
        line = line.replace("\r", "");
        line = line.replace("\n", "");
        line = line.strip();
        submask = "";
        network = "";
        if line.find("/") != -1:
            submask = line.split('/')[1];
            network = line.split('/')[0];
        else:
            network = line;
        try:
            network = socket.getaddrinfo(network, None)[0][4][0];
        except:
            continue;
        if submask != "":
            line = network+"/"+submask;
        else:
            line = network;
        if line.find('/') != -1:
            submask = line.split('/')[1];
            network = line.split('/')[0];
            [a1, a2, a3, a4] = network.split('.');
            u32 = (int(a1)<< 24) + (int(a2) << 16) + (int(a3) << 8) + int(a4);
            u32 &= 0xffffffff << (32 - int(submask));
            for i in xrange(1, 2 ** (32 - int(submask)) - 1):
                b1 = (u32 + i) >> 24;
                b2 = ((u32 + i) >> 16 ) & 0xff;
                b3 = ((u32 + i) >> 8 ) & 0xff;
                b4 = ((u32 + i) >> 0 ) & 0xff;
                targets.add(str(b1) + '.' + str(b2) + '.' + str(b3) + '.' + str(b4) );
        else:
            targets.add(line);
    tagetfile.close();
    return targets;


class autoLoadTargetThread(threading.Thread):
    def __init__(self, target_file, queue):
        self.target_file = target_file;
        self.queue = queue;
        self._stop = threading.Event();
        super(autoLoadTargetThread, self).__init__();

    def run(self):
        global targets
        global scanned_targets
        global targets_lock;
        global scanned_targets_lock;
        while True:
            if self._stop.isSet():
                break;
            time.sleep(60 * 15);
            mytarget = read_targets(self.target_file)
            targets_lock.acquire();
            scanned_targets_lock.acquire();
            for ip in mytarget:
                if ip not in targets and ip not in scanned_targets:
                    targets.add(ip);                    
                    self.queue.put(ip);
            scanned_targets_lock.release();
            targets_lock.release();
    
    def stop(self):
        self._stop.set()
        
def init():
    global logfile;
    global round;
    global logfile_service;
    print "start round %s scan"%(str(round))
    logfile = "log/"+datetime.datetime.now().strftime("%Y-%m-%d-%H-%M-%S") + ".txt"
    logfile_service = "logservice/"+datetime.datetime.now().strftime("%Y-%m-%d-%H-%M-%S") + ".txt"


def destroy():
    global round;
    print "end round %s scan"%(str(round))
    round += 1;

def _cliParser(args):
    parser = optparse.OptionParser();
    parser.add_option("-m", "--mode", action="store", type="string", help="specific mode to run", dest="mode", default="scanner");
    parser.add_option("-r",           action="store", type="string",help="which poc to load", dest="poc");
    parser.add_option( "--host", action="store", type="string", help="which host to test", dest="host");
    parser.add_option("-p", "--port", action="store", type="int", help="which port to test", dest="port");
    parser.add_option("-f", "--file", action="store", type="string", help="test file formats as HOST PORT SERVICE SERVICE_VERSION ", dest="testfile");
    parser.add_option("-v", "--version", action="store", type="string", help="use when mode is debugger identify service_version", dest="service_version");
    parser.add_option("-t", "--type", action="store", type="string", help="use when mode is debugger identify service_type ", dest="service_type");
    opts, args = parser.parse_args(args[1:])
    return opts, args


class batchThread(threading.Thread):
    def __init__(self, infodict, semp):
        self.infodict = infodict;
        self.semp = semp;
        super(batchThread, self).__init__();


    def run(self):
        for key in pocs.keys():
            localtmp = self.infodict.copy();
            if pocs.get(key)._match(localtmp):
                logToScreen("info",  "%s:executing %s for %s"%(self.name, key, localtmp["ip"]));
            ret = pocs.get(key).execute(localtmp);
            if ret.get("success") and ret["success"] == True:
                logToScreen("critical", key)
                logToScreen("critical", ret)
            if pocs.get(key)._match(localtmp):
                logToScreen("info",  "%s:execut %s for %s done"%(self.name, key, localtmp["ip"]));
        self.semp.release();




def main():
    global con_semp;
    global pro_semp;
    global scanned_targets;
    global targets;
    global targets_lock;
    banner();
    opt, baseurl = _cliParser(sys.argv);
    if opt.mode == "scanner":
        while True:
            loadpoc();
            producers = [];
            consumers = [];
            init();
            targets_lock.acquire();
            scan_queue = Queue();
            targets = read_targets("out.txt");
            for ip in targets:
                scan_queue.put(ip);
            targets_lock.release();
            autotarget = autoLoadTargetThread("out.txt", scan_queue)
            autopoc = autoLoadPocThread();
            autopoc.start();
            autotarget.start();
            while not scan_queue.empty():
                ip = scan_queue.get()
                def pro_callback():
                    con = consumerThread(reporter,con_semp );
                    consumers.append(con);
                    con.start();

                pro_semp.acquire();
                pro = produceThread(pro_semp, ip, con_semp, pro_callback);
                pro.start();
                producers.append(pro);
                
                scanned_targets_lock.acquire();
                scanned_targets.add(ip);
                scanned_targets_lock.release();
                scan_queue.task_done();
            autopoc.stop()
            autotarget.stop();
            autopoc.join();
            autotarget.join();
            for pro in producers:
                pro.join();
            for con in consumers:
                con.join();
            destroy();
    elif opt.mode == "debuger":
        from lib.classloader import loadspecpoc
        loadspecpoc(opt.poc);
        for value in pocs.values():
            ret = value.execute({"ip":opt.host, "port":opt.port, "service_type":opt.service_type, "service_version":opt.service_version});
            print ret;
    elif opt.mode == "batch":
        from lib.classloader import loadpocBydir;
        from lib.classloader import loadspecpoc
        loadspecpoc(opt.poc) if opt.poc.endswith(".py") else  loadpocBydir(opt.poc);
        localsemp = threading.Semaphore(100);
        threads = [];
        with open(opt.testfile) as myfile:
            for line in myfile.readlines():
                line = line.strip();
                line = line.replace("\r", "")
                line = line.replace("\n", "")
                tmp = {};
                tmp["ip"] = line.split()[0];
                tmp["port"] = line.split()[1];
                tmp["service_type"] = line.split()[3];
                tmp["service_version"] = " ".join(line.split()[4:]) if len(line.split()) > 4 else "";
                t = batchThread(tmp, localsemp);
                threads.append(t);
                localsemp.acquire();
                t.start();
        for t in threads:
            t.join();
if __name__ == "__main__":
    main();
